from yolov5 import YOLOv5
import cv2
import pickle as pk
import pdb
from tqdm import tqdm
import json
import os.path as osp
from scipy.io import loadmat
import torch
import argparse

def pickle(data, file_path):
    with open(file_path, "wb") as f:
        pk.dump(data, f, pk.HIGHEST_PROTOCOL)
def unpickle(file_path):
    with open(file_path, "rb") as f:
        data = pk.load(f)
    return data

def yolo_ssm_img(anno,yolo,root):
    print('start')
    det_bboxes_dic = {'person':[]}
    all_boxes=[]
    box_small6=[]
    box_smallhw15=[]
    box_smallh100=[]
    box_all_three=[]

    for i in tqdm(range(len(anno))):
        img_path=root+anno[i]
        src= cv2.imread(img_path)
        results = yolo.predict(img_path,size=640)
        # results.show()
        for result in results.xyxy[0]:
            result= result.cpu().numpy().tolist()
            x1,y1,x2,y2,score,cls, = result[:]
            x1,y1,x2,y2 = int(x1),int(y1),int(x2),int(y2)
            w,h=x2-x1,y2-y1
            if cls!=0:
                continue
            #筛选分数高>0.5的box,top10，h/w>1, h >50,w>50
            bb=[x1, y1, x2, y2,score,img_path]
            all_boxes.append(bb)

            if score<0.6: # h/w<1.5 or h<100 or
                box_small6.append(bb)
            if h/w<1.5:
                box_smallhw15.append(bb)
            if h<100:
                box_smallh100.append(bb)
            if h/w<1.5 or h<100 or score<0.6:
                box_all_three.append(bb)

            det_bboxes_dic['person'].append(
                                            [x1, y1, x2, y2,score,img_path]
            )
            cv2.rectangle(src,(x1,y1),(x2,y2),(0,0,255),2)
            cv2.putText(src, str(score),(x1,y1),1,1,(255,0,0),2)

        cv2.imshow('d',src)
        cv2.waitKey()
        # print(det_bboxes_dic)
        # print(len(det_bboxes_dic['person']))

    boxes_sum = [all_boxes,box_small6,box_smallhw15,box_smallh100,box_all_three]
    name = ['all_boxes','box_small6','box_smallhw15','box_smallh100','box_all_three']

    for i,n in zip(boxes_sum,name):
        print(n,len(i))
        pickle(i,'../../yolobox_analy/'+n+'.pkl')

    return det_bboxes_dic

def get_ssm_train_img_list(root='/home/cv7609/zjh/ps_raw/ssm'):
    # ssm= CUHKSYSU(root='/home2/0dataset/ps_raw/ssm',transforms=None,split='train')

    gallery_imgs = loadmat(osp.join(root, "annotation", "pool.mat"))
    gallery_imgs = gallery_imgs["pool"].squeeze()
    gallery_imgs = [str(a[0]) for a in gallery_imgs]
    # all images
    all_imgs = loadmat(osp.join(root, "annotation", "Images.mat"))
    all_imgs = all_imgs["Img"].squeeze()
    all_imgs = [str(a[0][0]) for a in all_imgs]
    # training images = all images - gallery images
    training_imgs = sorted(list(set(all_imgs) - set(gallery_imgs)))
    return training_imgs


if __name__=='__nmain__':
    root = '/home/cv7609/zjh/ps_raw/ssm'
    image_root = '/home/cv7609/zjh/ps_raw/ssm/Image/SSM/'
    train_imgs=get_ssm_train_img_list(root)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    yolo = YOLOv5('yolov5s.pt','cpu')
    #get yolo_det_ssm_json
    det_bboxes_dic=yolo_ssm_img(train_imgs,yolo,image_root)
    json_str = json.dumps(det_bboxes_dic,indent=4)
    json_name='../ssm_train_yolov5s_det_0.9.json'
    with open(json_name,'w') as json_file:
        json_file.write(json_str)

if __name__=='__nmain__':
    all_boxes = []
    box_small6 = []
    box_smallhw15 = []
    box_smallh100 = []
    box_all_three = []

    boxes_sum = [all_boxes, box_small6, box_smallhw15, box_smallh100, box_all_three]
    name = ['all_boxes', 'box_small6', 'box_smallhw15', 'box_smallh100', 'box_all_three']
    for i,n in enumerate(name):
        boxes_sum[i]=unpickle('../../yolobox_analy/' + n + '.pkl')
        print(n,len(boxes_sum[i]))

    # for init model > 0.6 > 1.5 person
    boxes_bigger6_bigger15=unpickle('../../yolobox_analy/' + 'boxes_bigger6_bigger15' + '.pkl')
    print('boxes_bigger6_bigger15:',len(boxes_bigger6_bigger15))

    filter=False
    if filter:
        boxes_bigger6_bigger15 = []
        for box in tqdm(boxes_sum[0]):
            if box in boxes_sum[1] or box in boxes_sum[2]:
                continue
            boxes_bigger6_bigger15.append(box)
        print(len(boxes_bigger6_bigger15))
        pickle(boxes_bigger6_bigger15, '../../yolobox_analy/' + 'boxes_bigger6_bigger15' + '.pkl')

if __name__=='__nmain__':
    im_p = '../s14859.jpg'
    src = cv2.imread(im_p)
    yolo = YOLOv5('yolov5s.pt', 'cpu')
    results = yolo.predict(im_p, size=640)
    # results.show()
    for result in results.xyxy[0]:
        result = result.cpu().numpy().tolist()
        x1, y1, x2, y2, score, cls, = result[:]
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        w, h = x2 - x1, y2 - y1
        cv2.rectangle(src, (x1, y1), (x2, y2), (0, 0, 255), 2)
    cv2.imshow('d',src)
    cv2.waitKey()

##------------------------------------------yolo------------------------------
if __name__=='__main__':
    parser = argparse.ArgumentParser(description="Train a person search network.")
    parser.add_argument("--score", default=0.3, type=float)
    show = False

    args = parser.parse_args()

    # image_root = '/home/cv7609/zjh/ps_raw/ssm/Image/SSM/'
    root = '/home/henry/file/0dataset/ps_raw/ssm/Image/SSM/'

    train_imgs = get_ssm_train_img_list('/home/henry/file/0dataset/ps_raw/ssm')
    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
    yolo = YOLOv5('yolov5l.pt', 'cuda:1')

    # get yolo_det_ssm_json
    det_bboxes_dic ={'person':[]}
    pids=0
    for i in tqdm(range(len(train_imgs))):
        img_path=root+train_imgs[i]
        src= cv2.imread(img_path)
        results = yolo.predict(img_path,size=640)
        # results.show()
        for result in results.xyxy[0]:
            result= result.cpu().numpy().tolist()
            x1,y1,x2,y2,score,cls, = result[:]
            x1,y1,x2,y2 = int(x1),int(y1),int(x2),int(y2)
            w,h=x2-x1,y2-y1
            if cls!=0 or score<args.score:
                continue
            #筛选分数高>0.5的box,top10，h/w>1, h >50,w>50
            # bb=[x1, y1, x2, y2,score,img_path]
            det_bboxes_dic['person'].append(
                [x1, y1, x2, y2, score, img_path]
            )
            pids+=1
            cv2.rectangle(src, (x1, y1), (x2, y2), (0, 0, 255), 2)
            cv2.putText(src, '{:.2f}'.format(score), (x1, y1), 1, 1, (255, 0, 0), 2)
        if show:
            cv2.imshow('d',src)
            cv2.waitKey()

    print(pids)
    json_str = json.dumps(det_bboxes_dic, indent=4)
    json_name = 'extra_ssm_yolo_'+str(args.score)+'_'+str(pids)+'.json'
    with open(json_name, 'w') as json_file:
        json_file.write(json_str)
